import datetime
import os
import pandas as pd
import Core.Config as Config
import Core.Gadget as Gadget


def calc_range_return(database, datetime1, datetime2):
    # 方便以后搜索到此文件
    pass


def rename_suffix(df, fields, suffix, inplace=False):
    rename_field = {}
    for field in fields:
        rename_field[field] = field + "_" + suffix
    df_rename = df.rename(columns=rename_field, inplace=inplace)
    return df_rename


def calc_stock_range_return(database, datetime1, datetime2):
    # 已经上市
    df_instrument = database.GetDataFrame("financial_data", "stock_instrument", projection=["symbol","datetime1","datetime2"])
    df_instrument_listed = df_instrument[(df_instrument["datetime1"] <= datetime1) & (df_instrument["datetime2"] >= datetime1)]

    # 读取价格收益
    trading_date_t1 = Gadget.find_recent_trading_day(database, datetime1, prev_next="prev")
    df_price_t1 = database.GetDataFrame("financial_data", "stock_daily_bar", filter=[("date", trading_date_t1)], projection=["symbol", "bclose","trade_status"])
    df_price_t1 = rename_suffix(df_price_t1, fields=["bclose","trade_status"], suffix="t1")

    #
    trading_date_t2 = Gadget.find_recent_trading_day(database, datetime2, prev_next="prev")
    df_price_t2 = database.GetDataFrame("financial_data", "stock_daily_bar", filter=[("date", trading_date_t2)], projection=["symbol", "bclose","trade_status"])
    df_price_t2 = rename_suffix(df_price_t2, fields=["bclose","trade_status"], suffix="t2")

    #
    df = pd.merge(df_price_t1, df_price_t2, how="inner", on="symbol")
    df = pd.merge(df_instrument_listed, df, how="left", on="symbol")
    df["range_return"] = df["bclose_t2"] / df["bclose_t1"] - 1
    # df.sort_values(by="range_return", inplace=True)
    df.dropna(subset=["range_return", "trade_status_t2"], inplace=True)
    # df_check = df[df["trade_status_t2"].isnull()]

    #
    def check_not_trade(row):
        if "停牌" in row["trade_status_t2"] and row["range_return"] == 0:
            return "not_trade"
        else:
            return "trade"

    #
    df["not_trade"] = df.apply(lambda x: check_not_trade(x), axis=1)
    df = df[~df["not_trade"].str.contains("not_trade")].copy()
    # print("range return", trading_date_t1, "to", trading_date_t2)
    #
    return df


# 不是迭代关系，两个均可用
def calc_stock_range_return_v2(database, datetime1, datetime2):

    real_datetime1 = Gadget.find_recent_date_in_table(database, "financial_data", "stock_dailybar", datetime1)
    real_datetime2 = Gadget.find_recent_date_in_table(database, "financial_data", "stock_dailybar", datetime2)

    df_1 = database.GetDataFrame("financial_data", "stock_dailybar", filter=[("date", real_datetime1)], projection=["symbol","bclose"])
    df_2 = database.GetDataFrame("financial_data", "stock_dailybar", filter=[("date", real_datetime2)], projection=["symbol","bclose"])
    #
    df_1.rename(columns={"bclose": "bclose_1"}, inplace=True)
    df_2.rename(columns={"bclose": "bclose_2"}, inplace=True)

    df = database.GetDataFrame("financial_data", "stock_instrument", projection=["symbol", "description", "industry"])

    df = pd.merge(df, df_1, on="symbol")
    df = pd.merge(df, df_2, on="symbol")
    #
    df["range_return"] = df["bclose_2"] / df["bclose_1"] - 1
    # df.dropna(subset=["return"], inplace=True)
    # df.sort_values(by="return", inplace=True, ascending=True)

    print("range return", real_datetime1, "to", real_datetime2)

    #
    return df[["symbol", "range_return"]].copy()